import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0, 1'

import csv
import pickle

import numpy as np
import tensorflow as tf


@tf.function(jit_compile=False)
def run(d):
    return tf.sets.intersection(**d)


@tf.function(jit_compile=True)
def run2(d):
    return tf.sets.intersection(**d)


if __name__ == "__main__":
    input_path = "/home/ubuntu/Ascend/model_inference/tf_source/optimize_inputs_1/tf.sets.intersection/tf.sets.intersection.pickle"
    compare_result_path = "/home/ubuntu/Ascend/model_inference/tf_results/optimize/tf.sets.intersection/20230316_122146/tf.sets.intersection.csv"
    source_result_path = "/home/ubuntu/Ascend/model_inference/tf_results/optimize/tf.sets.intersection/20230316_122146/tf.sets.intersection_source.pickle"
    follow_result_path = "/home/ubuntu/Ascend/model_inference/tf_results/optimize/tf.sets.intersection/20230316_122146/tf.sets.intersection_follow.pickle"

    with open(input_path, "rb") as f:
        input_dict_list = pickle.load(f)

    source_results = []
    follow_results = []

    for input_dict in input_dict_list:
        try:
            a = run(input_dict)
            print("output without XLA:")
            print(a)
            source_results.append(a.numpy())
        except Exception as e:
            print("[Not XLA] Exception occurred~:\n", e)
            source_results.append("Error")

        try:
            b = run2(input_dict)
            print("output with XLA:")
            print(b)
            follow_results.append(b.numpy())
        except Exception as e:
            print("[XLA] Exception occurred~:\n", e)
            follow_results.append("Error")

    with open(source_result_path, "wb+") as f:
        pickle.dump(source_results, f)

    with open(follow_result_path, "wb+") as f:
        pickle.dump(follow_results, f)

    compare_results = []
    for i in range(len(source_results)):
        a = source_results[i]
        b = follow_results[i]
        if type(a) is not np.ndarray and type(b) is not np.ndarray:
            if type(a) == str and a == "Error":
                compare_results.append([i, "Source Error"])
            elif type(b) == str and b == "Error":
                compare_results.append([i, "Follow Error"])
            else:
                if a == b:
                    compare_results.append([i, "Not Tensor Consistent"])
                else:
                    compare_results.append([i, "Not Tensor Inconsistent"])
        elif type(a) is not np.ndarray or type(b) is not np.ndarray:
            compare_results.append([i, "Inconsistent"])
        elif a.size == 0 and b.size == 0:
            compare_results.append([i, "Consistent"])
        elif a.size != b.size:
            compare_results.append([i, "Inconsistent", a.shape, b.shape])
        elif np.allclose(a, b, atol=1.e-2, rtol=1.e-5, equal_nan=True):
            compare_results.append([i, "Consistent", a.shape, b.shape, a.dtype, b.dtype])
        else:
            compare_results.append([i, "Inconsistent", a.shape, b.shape, a.dtype, b.dtype])
    with open(compare_result_path, "w+", newline="") as f:
        w = csv.writer(f)
        w.writerows(compare_results)
